#include "matmul.h"


template<typename T>
__global__ void cuda_matmul(T* A,T* B,T* C,int rows,int cols,int total_num){
    int idx = blockIdx.x*blockDim.x + threadIdx.x;
    if(idx<total_num){
        int row_idx = idx/cols;
        int col_idx = idx%cols;
        T temp = 0;
        for(int k=0;k<cols;++k){
            temp += A[row_idx*cols+k] * B[col_idx + k*rows];
        }
        C[row_idx*cols+col_idx] = temp;
    }
}


template<>
__global__ void cuda_matmul(half* A,half* B,half* C,int rows,int cols,int total_num){
    int idx = blockIdx.x*blockDim.x + threadIdx.x;
    if(idx<total_num){
        int row_idx = idx/cols;
        int col_idx = idx%cols;
        half temp = 0;
        for(int k=0;k<cols;++k){
            temp = __hadd(temp,__hmul(A[row_idx*cols+k] , B[col_idx + k*rows]));
        }
        C[row_idx*cols+col_idx] = temp;
    }
}



template<typename T>
void launch_cuda_matmul(T* A,T* B,T* C,int rows,int cols,int blocksize,int total_nums){
    int gridesize = (total_nums+blocksize -1)/blocksize;
    cuda_matmul<T><<<gridesize,blocksize>>>(A,B,C,rows,cols,total_nums);
}


template<typename T>
void launch_cpu_matmul(T* A,T* B,T* C,int rows,int cols){
    for(int row_idx=0;row_idx<rows;++row_idx){
        for(int col_idx=0;col_idx<cols;++col_idx){
            T temp = 0;
            for(int k=0;k<cols;++k){
                temp += A[row_idx*cols+k] * B[col_idx+k*rows];
            }
            C[row_idx*rows+col_idx] = temp;
        }
    }

}

template<>
void launch_cpu_matmul(half* A,half* B,half* C,int rows,int cols){
    for(int row_idx=0;row_idx<rows;++row_idx){
        for(int col_idx=0;col_idx<cols;++col_idx){
            float temp = 0;
            for(int k=0;k<cols;++k){
                // temp = __hadd(temp,__hmul(A[row_idx*cols+k] , B[col_idx+k*rows]));
                temp  += __half2float(A[row_idx*cols+k]) * __half2float(B[col_idx+k*rows]);
            }
            C[row_idx*rows+col_idx] = __float2half(temp);
        }
    }

}

template void launch_cuda_matmul<float>(float*,float*,float*,int,int,int,int);
template void launch_cuda_matmul<half>(half*,half*,half*,int,int,int,int);
template void launch_cuda_matmul<uint8_t>(uint8_t*,uint8_t*,uint8_t*,int,int,int,int);

template void launch_cpu_matmul<float>(float*,float*,float*,int,int);
template void launch_cpu_matmul<half>(half*,half*,half*,int,int);
template void launch_cpu_matmul<uint8_t>(uint8_t*,uint8_t*,uint8_t*,int,int);